{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c0e94770-1620-4693-ac3e-a9d59408f2c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.interactiveshell import InteractiveShell\n",
    "InteractiveShell.ast_node_interactivity = 'all' # default is ‘last_expr’\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "36d90bbc-4da9-404c-b621-11deae79c0eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "import json\n",
    "import os\n",
    "\n",
    "import ct_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2a090523-6d52-41d5-9703-2ae8d95962bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "list_images = os.listdir('test_images/test_images')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "279c9e45-1d6e-4e9b-9c38-f7d220d67336",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "31"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(list_images)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45e3f539-1315-4ebc-97e3-be25ae441e5d",
   "metadata": {},
   "source": [
    "Command used to run the detection script in the YOLOv5 repo: \n",
    "    \n",
    "```\n",
    "python detect.py --weights \"/Users/SiyuYang/OneDrive/CameraTraps/mdv5/yolov5_checkpoints/camonly_mosaic_xlarge_dist_5a_last.pt\" --source /Users/SiyuYang/Source/repos/Microsoft/CameraTraps/test_images/test_images --imgsz 1280 --conf-thres 0.01 --device cpu --save-txt --save-conf --project mdv5 --name detect_output\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ad4e2f4b-b6eb-4ec8-b2cc-64befca5ecc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "detect_script_output_dir = '/Users/SiyuYang/Source/repos/my_github/yolov5/mdv5/detect_output/labels'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b60dc2b7-2f18-456a-b545-16830f6bceba",
   "metadata": {},
   "outputs": [],
   "source": [
    "yolo_cat_map = {\n",
    "    0: 1,\n",
    "    1: 2,\n",
    "    2: 3\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "aeff9627-1a0d-4314-bfe9-31e2ab8e589e",
   "metadata": {},
   "outputs": [],
   "source": [
    "images_entries = []\n",
    "\n",
    "for image_fn in list_images:\n",
    "    image_name, ext = os.path.splitext(image_fn)\n",
    "    \n",
    "    label_fn = image_name + '.txt'\n",
    "    label_path = os.path.join(detect_script_output_dir, label_fn)\n",
    "        \n",
    "    detections = []\n",
    "    max_conf = 0.0\n",
    "    \n",
    "    if os.path.exists(label_path):\n",
    "        with open(label_path, newline='') as f:\n",
    "            reader = csv.reader(f, delimiter=' ')\n",
    "            for row in reader:\n",
    "                category = yolo_cat_map[int(row[0])]\n",
    "\n",
    "                api_box = ct_utils.convert_yolo_to_xywh([float(row[1]), float(row[2]), float(row[3]), float(row[4])])\n",
    "\n",
    "                conf = ct_utils.truncate_float(float(row[5]), precision=4) # output precision is 6 decimal places\n",
    "                max_conf = max(max_conf, conf)\n",
    "\n",
    "                detections.append({\n",
    "                    'category': str(category),\n",
    "                    'conf': conf,\n",
    "                    'bbox': ct_utils.truncate_float_array(api_box, precision=4)\n",
    "                })\n",
    "            \n",
    "    images_entries.append({\n",
    "        'file': image_fn,\n",
    "        'max_detection_conf': max_conf,\n",
    "        'detections': detections\n",
    "    })\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "8da4c5f9-89ca-413a-b182-943f2edc0d04",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_content = {\n",
    "    'info': {\n",
    "        'detector': 'megadetector_v5_camonly_mosaic_xlarge_dist_5a_last',\n",
    "        'format_version': '1.1'\n",
    "    },\n",
    "    'detection_categories': {\n",
    "        '1': 'animal',\n",
    "        '2': 'person',\n",
    "        '3': 'vehicle'\n",
    "    },\n",
    "    'images': images_entries\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "aba09987-25ba-4ec1-a6dd-9361e4169bfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('test_images/expected_output_DRAFT.json', 'w') as f:\n",
    "    json.dump(output_content, f, indent=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63f3e5c0-0d19-4a59-9083-106e118e2db0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
